import numpy as np
from matplotlib import pyplot as plt
from sklearn.cross_decomposition import PLSRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

from input import input_data


class PLS_usage_package():
    def __init__(self,n_components=3):# 选择一个 n_components 值，这个值影响拟合效果
        self.model = None
        self.X_raw = None
        self.Y = None
        self.pls = PLSRegression(n_components=n_components)  # 选择一个 n_components 值，这个值影响拟合效果

    def input_data(self,X_raw,Y):
        '''
        输入数据
        :param X_raw: np.array([intensity1, intensity2, ...])光强值二维数组n个光强值*m个样品
        :param Y: np.array([concentration1, concentration2, ...])浓度值一维数组m个样品*1
        :return: 无
        '''
        self.X_raw = X_raw
        #print(self.X_raw)
        self.Y = Y
        #print(self.Y.shape)

    def create_PLS_model(self):
        '''
        创建模型，注意，创建前一定先要input数据集进去
        :return:None
        '''
        X = self.X_raw
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, self.Y, test_size=0.2, random_state=42)
        # 创建并训练 PLSRegression 模型
        self.pls.fit(X_train, y_train)
        # 预测
        y_pred = self.pls.predict(X_test)
        # plt.plot(X_test[2],label='pre')
        # plt.show()
        # 评估模型
        mse = mean_squared_error(y_test, y_pred)
        print(f"Mean Squared Error: {mse}")

    def inference(self,data):
        '''
        根据当前模型推演预测浓度
        :param data: np.array 1xm，或者nxm的矩阵，每个矩阵里放着光强数列
        :return: np.array nx1的矩阵，每一行放着预测好的浓度
        '''
        #print(data.shape)
        y_pred = self.pls.predict(data)
        print(f"预测的浓度为{y_pred}")
        return y_pred





if __name__ == "__main__":
    pls = PLS_usage_package()
    matrix,percentages = input_data()
    #print(matrix.shape)
    percentages = np.array(percentages)
    # 假设我们有100个样本，每个样本在200个不同的波长点进行测量
    # num_samples = 100
    # 随机生成对应的光强度值（光强度0-1之间）
    # intensities = np.random.rand(200, num_samples)
    # 打印一个样本以验证数据格式
    #print(X_raw[0])
    # Y = np.random.rand(num_samples, 1)
    pls.input_data(matrix,percentages)
    pls.create_PLS_model()
    #选取一个文件进行浓度预测
    inference_data,_ = input_data(file_paths="F:\\pythonProject\\data_processing_platform\\resources\\data")
    #print(inference_data.shape)
    pls.inference(inference_data)